import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)

import numpy as np
import matplotlib.colors as mcolors
from matplotlib import pyplot as plt

from utils import *

X_RANGE = (-40, 40)
Y_RANGE = (0, 90)

def filter_data(X, Y, Z, x_range, y_range, azi, ele, err):
    mask = (Z != 0) & \
        (X >= x_range[0]) & (X <= x_range[1]) & \
        (Y >= y_range[0]) & (Y <= y_range[1]) & ((azi > 0.9) & (ele > 0.9)) & (err < 0.2 * np.sqrt(X**2 + Y**2))
    return mask

def filter_and_visualize(data, output_folder, x_range=X_RANGE, y_range=Y_RANGE):
    utc_time = data.get('utcTimeChar', ["Unknown"])[0].replace(':', '_')

    os.makedirs(output_folder, exist_ok=True)
    output_file = os.path.join(output_folder, f"{utc_time}.png")

    if os.path.exists(output_file):
        return
    
    azi_map = data.get('aziLinearlityMap', None)
    ele_map = data.get('eleLinearlityMap', None)
    error_map = data.get('errorMap', None)
    target_x = data.get('targetX', None)
    target_y = data.get('targetY', None)
    target_z = data.get('targetZ', None)
    dn = data.get('DN', None)
    
    if not all([azi_map is not None, ele_map is not None, error_map is not None, target_x is not None, target_y is not None, target_z is not None, dn is not None]):
        raise ".mat文件中缺少数据"

    X = target_x.flatten()
    Y = target_y.flatten()
    Z = target_z.flatten()
    DN = dn.flatten()
    azi = azi_map.flatten()
    ele = ele_map.flatten()

    mask = filter_data(X, Y, Z, x_range, y_range, azi, ele, err)

    X = X[mask]
    Y = Y[mask]
    Z = Z[mask]
    DN = DN[mask]
    s = s[mask]
    azi = azi[mask]
    ele = ele[mask]

    hue = 0.5 * (Z - Z.min()) / (Z.max() - Z.min())
    value = (s - s.min()) / (s.max() - s.min())
    saturation = 1 - (DN - DN.min()) / (DN.max() - DN.min())
    
    hsv = np.stack((hue, saturation, value), axis=1)
    rgb = mcolors.hsv_to_rgb(hsv)

    fig, ax = plt.subplots(figsize=(10, 10), dpi=1024//10, facecolor='black')
    ax.set_facecolor('blue')
    ax.set_xlim(x_range)
    ax.set_ylim(y_range)
    ax.set_axis_off()
    ax.scatter(X, Y, c=rgb, s=1, marker='.')

    fig.savefig(output_file, dpi=1024//10, facecolor=fig.get_facecolor())
    plt.close()

def main():
    folder_path = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR/each"
    
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".mat"):
                file_path = os.path.join(root, file)
                print(f"Processing {file_path}")

                mat_data = read_mat_file(file_path)
                if mat_data:
                    filter_and_visualize(mat_data, output_folder)

if __name__ == "__main__":
    main()
